from Neural_Overlapping_Quantization.dataset_partition_1 import kmeans_multiple_overlap_boundary, \
    kmeans_multiple, knn
from Neural_Overlapping_Quantization.util import dir_io
import time
import numpy as np
import copy
import os


def dataset_partition(base, base_base_gnd, config):
    dir_io.mkdir(config['program_train_para_dir'])
    start_time = time.time()
    partition_method = factory(config['dataset_partition']['type'])
    for i in range(config['n_classifier']):
        save_dir = '%s/Classifier_%d/' % (config['program_train_para_dir'], i)
        os.mkdir(save_dir)
        save_dir = '%s/Classifier_%d/dataset_partition' % (config['program_train_para_dir'], i)
        os.mkdir(save_dir)
    print("start dataset partition")
    label_l, intermediate = partition_method(base, base_base_gnd, config)
    print("end dataset partition")
    label_map_l = []
    for i, label in enumerate(label_l):
        label_map, n_point_label = get_label_map(label, config['n_cluster'])
        intermediate_name = "n_point_label_%d" % i
        intermediate['dataset_partition'][i][intermediate_name] = n_point_label
        label_map_l.append(label_map)
    end_time = time.time()

    intermediate['total_dataset_partition_time'] = end_time - start_time
    return (label_l, label_map_l), intermediate


# the function partition the base according to the number of cluster
# the labels should be the array of numpy
def get_label_map(labels, n_cluster):
    label_map = []
    n_point_label = []
    for cluster_i in range(n_cluster):
        base_idx_i = np.argwhere(labels == cluster_i)[:, 0]
        label_map.append(base_idx_i)
        n_point_label.append(len(base_idx_i))
    return label_map, n_point_label


def factory(_type):
    if _type == 'kmeans_multiple_overlap_boundary':
        return kmeans_multiple_overlap_boundary.dataset_partition
    elif _type == 'kmeans_multiple':
        return kmeans_multiple.dataset_partition
    elif _type == 'knn':
        return knn.dataset_partition
    raise Exception('do not support the type of partition')
